#include <tofino/constants.p4>
#include <tofino/intrinsic_metadata.p4>
#include <tofino/primitives.p4>
#include <tofino/pktgen_headers.p4>
#include <tofino/stateful_alu_blackbox.p4>
#include <tofino/wred_blackbox.p4>

#include "includes/headers.p4"
#include "includes/parser.p4"
#include "includes/macro.p4"
#include "includes/tofino.p4"

field_list flow {
    ipv4.srcip;
    ipv4.dstip;
    ipv4.proto;
    tcp.srcport;
    tcp.dstport;
}

// for the calculation of the index of the filter of the ingress sketch
field_list_calculation filter_idx_hash {
    input {
        flow;
    }
    algorithm: crc32;
    output_width: INGRESS_SKETCH_IDX_WIDTH;
}

// for the calculation of the index of the row 1 of the ingress sketch
field_list_calculation is_row1_idx_hash {
    input {
        flow;
    }
    algorithm: crc32_extend;
    output_width: INGRESS_SKETCH_IDX_WIDTH;
}

// for the calculation of the index of the row 2 of the ingress sketch
field_list_calculation is_row2_idx_hash {
    input {
        flow;
    }
    algorithm: crc32_lsb;
    output_width: INGRESS_SKETCH_IDX_WIDTH;
}

field_list_calculation fingerprint_hash {
    input {
        flow;
    }
    algorithm: identity;
    output_width: FINGERPRINT_WIDTH;
}

// metadata for measurement program
header_type measurement_meta_t {
    fields {
        fingerprint: FINGERPRINT_WIDTH; // fingerprint for 5-tuple;
        filter_predicate: 4; // output predicate of the filter;
		minbound: 32; // the minimum upper bound for the current flow
		npkts: 32; // the total number of packets transmitted by the switch
		prior_minbound: 32;
		prior_npkts: 32;
		es_timestamp_flag: 1; // if the bucket of ES is outdated.
		es_fingerprint_flag: 1; // if the fingerprint of the bucket of ES matches with the current packet
		deq_timestamp: 32;
		delay_flag: 32;
		count_flag: 32;
		minbound_diff: 32;
    }
}

metadata measurement_meta_t measurement_meta;


action set_egr(egress_spec) {
	modify_field(ig_intr_md_for_tm.ucast_egress_port, egress_spec);
}
 
action nop() {
}
 
@pragma stage 0
table forward {
	reads {
		ig_intr_md.ingress_port: exact;
	}
	actions {
		set_egr; nop;
	}
}

action calc_minbound_diff() {
	subtract(measurement_meta.minbound_diff, measurement_meta.minbound, \
		measurement_meta.prior_minbound);
}

action calc_count_flag() {
	subtract(measurement_meta.count_flag, COUNT_THRESH, measurement_meta.minbound_diff);
}

// action for transforming 5-tuple into 32-bit fingerprint;
action generate_fingerprint_action()
{
    modify_field_with_hash_based_offset(measurement_meta.fingerprint, 0, fingerprint_hash, 0x10000);
}

// table for transforming 5-tuple into 32-bit fingerprint;
@pragma stage 0
table generate_fingerprint_t
{
    actions {
        generate_fingerprint_action;
    }
    default_action: generate_fingerprint_action;
}

// register for the total number of packets
register npkts
{
	width: 32;
	instance_count: 1;
}

blackbox stateful_alu update_npkts
{
    reg: npkts;
    update_lo_1_value: register_lo + 1;
    output_value: alu_lo;
    output_dst: measurement_meta.npkts;
}

action update_npkts_action()
{
    update_npkts.execute_stateful_alu();
}

@pragma stage 0
table update_npkts_t
{
    actions {
		update_npkts_action;
    }
    default_action: update_npkts_action;
}

// register for the fingerprints and temporary counters of filter
register filter_fingerprint
{
    width: 32;
    instance_count: INGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_filter_fingerprint
{
    reg: filter_fingerprint;
    condition_lo: register_lo == measurement_meta.fingerprint;
    condition_hi: register_hi == 0;

    update_lo_1_predicate: condition_lo or condition_hi;
    update_lo_1_value: measurement_meta.fingerprint;

    update_hi_1_predicate: condition_lo and condition_hi;
    update_hi_1_value: register_hi + 1;
    update_hi_2_predicate: not condition_lo and not condition_hi;
    update_hi_2_value: register_hi - 1;

    output_value: predicate;
    output_dst: measurement_meta.filter_predicate;
}

action update_filter_fingerprint_action()
{
    update_filter_fingerprint.execute_stateful_alu_from_hash(filter_idx_hash);
}

@pragma stage 1
table update_filter_fingerprint_t
{
    actions {
		update_filter_fingerprint_action;
    }
    default_action: update_filter_fingerprint_action;
}

register filter_counter
{
    width: 32;
    instance_count: INGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_filter_counter
{
    reg: filter_counter;
    update_lo_1_value: register_lo + 1;
}

action update_filter_counter_action()
{
    update_filter_counter.execute_stateful_alu_from_hash(filter_idx_hash);
}

blackbox stateful_alu read_filter_counter
{
    reg: filter_counter;
    output_value: register_lo;
    output_dst: measurement_meta.minbound;
}

action read_filter_counter_action()
{
    read_filter_counter.execute_stateful_alu_from_hash(filter_idx_hash);
}

@pragma stage 2
table process_filter_counter_t
{
    reads {
        measurement_meta.filter_predicate: exact;
    }
    actions {
        update_filter_counter_action;
        read_filter_counter_action;
    }
	max_size: 4;
}

register is_row1
{
    width: 32;
    instance_count: INGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_is_row1
{
    reg: is_row1;
    condition_lo: register_lo < measurement_meta.minbound;

    update_lo_1_predicate: condition_lo;
    update_lo_1_value: register_lo + 1;

	output_predicate: condition_lo;
    output_value: alu_lo;
    output_dst: measurement_meta.minbound;
}

action update_is_row1_action()
{
    update_is_row1.execute_stateful_alu_from_hash(is_row1_idx_hash);
}

@pragma stage 3
table update_is_row1_t
{
    actions {
		update_is_row1_action;	
    }
    default_action: update_is_row1_action;
}

register is_row2
{
    width: 32;
    instance_count: INGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_is_row2
{
    reg: is_row2;
    condition_lo: register_lo < measurement_meta.minbound;

    update_lo_1_predicate: condition_lo;
    update_lo_1_value: register_lo + 1;

	output_predicate: condition_lo;
    output_value: alu_lo;
    output_dst: measurement_meta.minbound;
}

action update_is_row2_action()
{
    update_is_row2.execute_stateful_alu_from_hash(is_row2_idx_hash);
}

@pragma stage 4
table update_is_row2_t
{
    actions {
		update_is_row2_action;	
    }
    default_action: update_is_row2_action;
}

control ingress
{
	// stage 0
	apply(forward); 
	if(valid(tcp) or valid(udp)) {
		// stage 0
		apply(update_npkts_t); 
		apply(generate_fingerprint_t);
		// stage 1
		apply(update_filter_fingerprint_t);
		// stage 2
		apply(process_filter_counter_t);
		// stage 3
		apply(update_is_row1_t);
		// stage 4
		apply(update_is_row2_t);
	}
}


// the following are for the egress sketch
field_list_calculation es_row1_idx_hash {
    input {
        flow;
    }
    algorithm: crc32;
    output_width: EGRESS_SKETCH_IDX_WIDTH;
}

field_list_calculation es_row2_idx_hash {
    input {
        flow;
    }
    algorithm: crc32_extend;
    output_width: EGRESS_SKETCH_IDX_WIDTH;
}

field_list_calculation es_row3_idx_hash {
    input {
        flow;
    }
    algorithm: crc32_lsb;
    output_width: EGRESS_SKETCH_IDX_WIDTH;
}

field_list_calculation es_row4_idx_hash {
    input {
        flow;
    }
    algorithm: crc32_msb;
    output_width: EGRESS_SKETCH_IDX_WIDTH;
}

action get_deq_timestamp()
{
	add(measurement_meta.deq_timestamp, intrinsic_metadata.enq_timestamp, \
		intrinsic_metadata.deq_timedelta);
	subtract(measurement_meta.delay_flag, DELAY_THRESH, intrinsic_metadata.deq_timedelta);
}

@pragma stage 0
table get_deq_timestamp_t {
	actions {
		get_deq_timestamp;
	}
	default_action: get_deq_timestamp;
}

/*ES Row1 timestamp*/
register es_row1_timestamp
{
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row1_timestamp
{
	reg: es_row1_timestamp;
	condition_lo: register_lo <= intrinsic_metadata.enq_timestamp;

	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.deq_timestamp;

	output_value: predicate;
	output_dst: measurement_meta.es_timestamp_flag;
}

action update_es_row1_timestamp_action()
{
	update_es_row1_timestamp.execute_stateful_alu_from_hash(es_row1_idx_hash);
}

@pragma stage 1
table update_es_row1_timestamp_t 
{
	actions {
		update_es_row1_timestamp_action;
	}
	default_action: update_es_row1_timestamp_action;
}

/*ES Row1 MinBound*/
register es_row1_minbound {
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row1_minbound {
	reg: es_row1_minbound;
	condition_lo: measurement_meta.es_timestamp_flag == 2;
	
	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.minbound;

	output_value: register_lo;
	output_dst: measurement_meta.prior_minbound;

}

action update_es_row1_minbound_action() {
	update_es_row1_minbound.execute_stateful_alu_from_hash(es_row1_idx_hash);
}

@pragma stage 2
table update_es_row1_minbound_t {
	actions {
		update_es_row1_minbound_action;
	}
	default_action: update_es_row1_minbound_action;
}
/*ES Row1 fingerprint*/
register es_row1_fingerprint {
	width: 16;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row1_fingerprint {
	reg: es_row1_fingerprint;
	condition_lo: register_lo == measurement_meta.fingerprint;
	condition_hi: measurement_meta.es_timestamp_flag == 2;

	update_lo_1_predicate: condition_hi;
	update_lo_1_value: measurement_meta.fingerprint;

	output_value: predicate;
	output_dst: measurement_meta.es_fingerprint_flag;
}

action update_es_row1_fingerprint_action() {
	update_es_row1_fingerprint.execute_stateful_alu_from_hash(es_row1_idx_hash);
	calc_minbound_diff();
}

@pragma stage 3
table update_es_row1_fingerprint_t {
	actions {
		update_es_row1_fingerprint_action;
	}
	default_action: update_es_row1_fingerprint_action;
}


/*ES Row1 npkts*/
register es_row1_npkts {
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row1_npkts {
	reg: es_row1_npkts;
	condition_lo: measurement_meta.es_timestamp_flag == 2;
	
	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.npkts;

	output_value: register_lo;
	output_dst: measurement_meta.prior_npkts;
}

action update_es_row1_npkts_action() {
	update_es_row1_npkts.execute_stateful_alu_from_hash(es_row1_idx_hash);
	calc_count_flag();
}

@pragma stage 4
table update_es_row1_npkts_t {
	reads {
		measurement_meta.es_timestamp_flag: exact;
	}
	actions {
		update_es_row1_npkts_action; nop;
	}
	default_action: nop;
}

/*ES Row2 timestamp*/
register es_row2_timestamp
{
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row2_timestamp
{
	reg: es_row2_timestamp;
	condition_lo: register_lo <= intrinsic_metadata.enq_timestamp;

	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.deq_timestamp;

	output_value: predicate;
	output_dst: measurement_meta.es_timestamp_flag;
}

action update_es_row2_timestamp_action()
{
	update_es_row2_timestamp.execute_stateful_alu_from_hash(es_row2_idx_hash);
}

@pragma stage 5
table update_es_row2_timestamp_t 
{
	actions {
		update_es_row2_timestamp_action;
	}
	default_action: update_es_row2_timestamp_action;
}

/*ES Row2 MinBound*/
register es_row2_minbound {
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row2_minbound {
	reg: es_row2_minbound;
	condition_lo: measurement_meta.es_timestamp_flag == 2;
	
	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.minbound;

	output_value: register_lo;
	output_dst: measurement_meta.prior_minbound;

}

action update_es_row2_minbound_action() {
	update_es_row2_minbound.execute_stateful_alu_from_hash(es_row2_idx_hash);
}

@pragma stage 6
table update_es_row2_minbound_t {
	actions {
		update_es_row2_minbound_action;
	}
	default_action: update_es_row2_minbound_action;
}

/*ES Row2 fingerprint*/
register es_row2_fingerprint {
	width: 16;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row2_fingerprint {
	reg: es_row2_fingerprint;
	condition_lo: register_lo == measurement_meta.fingerprint;
	condition_hi: measurement_meta.es_timestamp_flag == 2;

	update_lo_1_predicate: condition_hi;
	update_lo_1_value: measurement_meta.fingerprint;

	output_value: predicate;
	output_dst: measurement_meta.es_fingerprint_flag;
}

action update_es_row2_fingerprint_action() {
	update_es_row2_fingerprint.execute_stateful_alu_from_hash(es_row2_idx_hash);
	calc_minbound_diff();
}

@pragma stage 7
table update_es_row2_fingerprint_t {
	actions {
		update_es_row2_fingerprint_action;
	}
	default_action: update_es_row2_fingerprint_action;
}


/*ES Row2 npkts*/
register es_row2_npkts {
	width: 32;
	instance_count: EGRESS_SKETCH_WIDTH;
}

blackbox stateful_alu update_es_row2_npkts {
	reg: es_row2_npkts;
	condition_lo: measurement_meta.es_timestamp_flag == 2;
	
	update_lo_1_predicate: condition_lo;
	update_lo_1_value: measurement_meta.npkts;

	output_value: register_lo;
	output_dst: measurement_meta.prior_npkts;
}

action update_es_row2_npkts_action() {
	update_es_row2_npkts.execute_stateful_alu_from_hash(es_row2_idx_hash);
	calc_count_flag();
}

@pragma stage 8
table update_es_row2_npkts_t {
	reads {
		measurement_meta.es_timestamp_flag: exact;
	}
	actions {
		update_es_row2_npkts_action; nop;
	}
	default_action: nop;
}

// process the tcp and udp headers
action record_ip_address() {
	add_header(perep);
	modify_field(perep.srcip, ipv4.srcip);
	modify_field(perep.dstip, ipv4.dstip);
	modify_field(perep.proto, ipv4.proto);
	modify_field(perep.enq_qdepth, intrinsic_metadata.enq_qdepth);
}

action process_tcp_action() {
	record_ip_address();
	modify_field(perep.srcport, tcp.srcport);
	modify_field(perep.dstport, tcp.dstport);
	subtract_from_field(ipv4.totalLen, 12);
	remove_header(tcp);
	add_header(udp);
}

action process_udp_action() {
	record_ip_address();
	modify_field(perep.srcport, udp.srcport);
	modify_field(perep.dstport, udp.dstport);
}

@pragma stage 10
table process_udp_t {
	actions {
		process_udp_action;
	}
	default_action: process_udp_action;
}

@pragma stage 10
table process_tcp_t {
	actions {
		process_tcp_action;
	}
	default_action: process_tcp_action;
}

@pragma stage 9
table process_udp_tcp_t {
	reads {
		udp: valid;
		tcp: valid;
		measurement_meta.delay_flag: ternary;
		measurement_meta.count_flag: ternary;
	}
	actions {
		process_udp_action;
		process_tcp_action;
		nop;
	}
	default_action: nop;
	max_size: 2;
}

// report the congestion maker
action report_congestion_maker_action() {
	modify_field(udp.dstport, UDP_PEREP);
	modify_field(udp.checksum, 0);
	modify_field(perep.deq_timedelta, intrinsic_metadata.deq_timedelta);
	modify_field(perep.minbound, measurement_meta.minbound);
	modify_field(perep.prior_minbound, measurement_meta.prior_minbound);
	modify_field(perep.npkts, measurement_meta.npkts);
	modify_field(perep.prior_npkts, measurement_meta.prior_npkts);
	add(ipv4.totalLen, ipv4.totalLen, 44);
	modify_field(ipv4.dstip, CONTROL_HOST_ADDR);
	modify_field(ig_intr_md_for_tm.ucast_egress_port, CONTROL_EGRESS_PORT);
}

@pragma stage 10
table report_congestion_maker_t {
	reads {
		measurement_meta.delay_flag: ternary;
		measurement_meta.count_flag: ternary;
	}
	actions {
		report_congestion_maker_action; nop;
	}
	default_action: nop;
}

control egress
{
	// stage 0
	apply(get_deq_timestamp_t);
	if(valid(udp) or valid(tcp)) {
		// stage 1
		apply(update_es_row1_timestamp_t);
		// stage 2
		apply(update_es_row1_minbound_t);
		// stage 3
		apply(update_es_row1_fingerprint_t);
		// stage 4
		apply(update_es_row1_npkts_t);
		if(1 == measurement_meta.es_fingerprint_flag) {
			// stage 5
			apply(update_es_row2_timestamp_t);
			// stage 6
			apply(update_es_row2_minbound_t);
			// stage 7
			apply(update_es_row2_fingerprint_t);
			// stage 8
			apply(update_es_row2_npkts_t);
		}
		if(8 == measurement_meta.es_fingerprint_flag) {
			// stage 9
			apply(process_udp_tcp_t);
			// stage 10
			apply(report_congestion_maker_t);
		}
	}
}
